{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9a47c35e-5fd7-48ca-b865-d07d76d3cfff",
   "metadata": {},
   "source": [
    "# L1: Vanilla Vector Search\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01acdf11-0273-4a1a-9b35-090bfeb44343",
   "metadata": {},
   "source": [
    "<p style=\"background-color:#fff6e4; padding:15px; border-width:3px; border-color:#f5ecda; border-style:solid; border-radius:6px\"> ⏳ <b>Note <code>(Kernel Starting)</code>:</b> This notebook takes about 30 seconds to be ready to use. You may start and watch the video while you wait.</p>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8062eca-e2de-4c79-ad94-9b4ad8e70d7e",
   "metadata": {
    "height": 64
   },
   "outputs": [],
   "source": [
    "# 警告控制\n",
    "# 忽略所有的警告信息，以免干扰程序输出\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b5cd7ff-ff61-4fc2-a246-b6ffeaa2869e",
   "metadata": {
    "height": 47
   },
   "outputs": [],
   "source": [
    "#!pip install datasets pandas openai pymongo pydantic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "835995e2-98e6-46cb-8050-9dce6d44f5c8",
   "metadata": {},
   "source": [
    "## Get API Keys\n",
    "In this classroom, the libraries and APIs have been already installed and set up for you.\n",
    "If you would like to run this code on your own machine, you will need to enter your own MONGO_URI and OPENAI_API_KEY keys in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32321449-c1ca-4851-95cd-8992418bdc1c",
   "metadata": {
    "height": 115
   },
   "outputs": [],
   "source": [
    "# 导入所需的库\n",
    "import os\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "\n",
    "# 读取本地的 .env 文件中的环境变量\n",
    "_ = load_dotenv(find_dotenv())\n",
    "\n",
    "# 获取环境变量 OPENAI_API_KEY 的值\n",
    "OPENAI_API_KEY = os.environ.get(\"OPENAI_API_KEY\")\n",
    "\n",
    "# 获取环境变量 MONGO_URI 的值\n",
    "MONGO_URI = os.environ.get(\"MONGO_URI\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f7e67f7-f17c-4d18-ba19-ec64d4d544d7",
   "metadata": {},
   "source": [
    "<p style=\"background-color:#fff6ff; padding:15px; border-width:3px; border-color:#efe6ef; border-style:solid; border-radius:6px\"> 💻 &nbsp; <b>Access <code>requirements.txt</code> file:</b> To access <code>requirements.txt</code> for this notebook, 1) click on the <em>\"File\"</em> option on the top menu of the notebook and then 2) click on <em>\"Open\"</em>. For more help, please see the <em>\"Appendix - Tips and Help\"</em> Lesson.</p>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b90f1121-a49e-45fc-af27-6da418c291cc",
   "metadata": {},
   "source": [
    "## 1.1 Data Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6104c624-b988-4ea6-9c79-f76907891507",
   "metadata": {
    "height": 251
   },
   "outputs": [],
   "source": [
    "# 1. 加载数据集\n",
    "from datasets import load_dataset\n",
    "import pandas as pd\n",
    "\n",
    "# 注意：确保在开发环境中有 Hugging Face 令牌 (HF_TOKEN)\n",
    "# 参考：https://huggingface.co/datasets/MongoDB/airbnb_embeddings\n",
    "# 注意：此数据集包含多个记录，每个记录代表一个 Airbnb 列表项\n",
    "# 注意：此数据集包含文本和图像嵌入，但本教程仅使用文本嵌入\n",
    "dataset = load_dataset(\"MongoDB/airbnb_embeddings\", streaming=True, split=\"train\")\n",
    "dataset = dataset.take(100)  # 取前100条数据\n",
    "\n",
    "# 将数据集转换为 pandas 数据框\n",
    "dataset_df = pd.DataFrame(dataset)\n",
    "\n",
    "# 显示前5条数据\n",
    "dataset_df.head(5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5194bdd7-f563-4bdb-bc49-b528beaf5543",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "# 打印数据框的列名\n",
    "print(\"Columns:\", dataset_df.columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de8bc707-c6bb-4ff4-b801-a3db598ae253",
   "metadata": {},
   "source": [
    "## 1.2 Document Modelling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab58b04e-510e-44ab-bde0-9ceefdbb80d8",
   "metadata": {
    "height": 64
   },
   "outputs": [],
   "source": [
    "# 导入所需的库\n",
    "from typing import List, Optional\n",
    "from pydantic import BaseModel, ValidationError\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d91abf6f-6272-4577-a6e6-65aa9d06887b",
   "metadata": {
    "height": 234
   },
   "outputs": [],
   "source": [
    "# 定义 Host 类，该类继承自 BaseModel，用于表示房东的相关信息\n",
    "class Host(BaseModel):\n",
    "    host_id: str  # 房东ID\n",
    "    host_url: str  # 房东主页URL\n",
    "    host_name: str  # 房东姓名\n",
    "    host_location: str  # 房东位置\n",
    "    host_about: str  # 关于房东的描述\n",
    "    host_response_time: Optional[str] = None  # 房东响应时间，可选\n",
    "    host_thumbnail_url: str  # 房东缩略图URL\n",
    "    host_picture_url: str  # 房东图片URL\n",
    "    host_response_rate: Optional[int] = None  # 房东响应率，可选\n",
    "    host_is_superhost: bool  # 是否为超级房东\n",
    "    host_has_profile_pic: bool  # 是否有个人资料图片\n",
    "    host_identity_verified: bool  # 身份是否已验证"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb99dd71-377e-4036-8651-36e7292078e1",
   "metadata": {
    "height": 217
   },
   "outputs": [],
   "source": [
    "# 定义 Location 类，用于表示位置相关信息\n",
    "class Location(BaseModel):\n",
    "    type: str  # 位置类型\n",
    "    coordinates: List[float]  # 坐标列表\n",
    "    is_location_exact: bool  # 位置是否准确\n",
    "\n",
    "# 定义 Address 类，用于表示地址相关信息\n",
    "class Address(BaseModel):\n",
    "    street: str  # 街道\n",
    "    government_area: str  # 政府区域\n",
    "    market: str  # 市场\n",
    "    country: str  # 国家\n",
    "    country_code: str  # 国家代码\n",
    "    location: Location  # 位置对象\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c647ebab-8e65-4011-8a11-994eefa7051f",
   "metadata": {
    "height": 132
   },
   "outputs": [],
   "source": [
    "# 定义 Review 类，用于表示评论相关信息\n",
    "class Review(BaseModel):\n",
    "    _id: str  # 评论ID\n",
    "    date: Optional[datetime] = None  # 评论日期，可选\n",
    "    listing_id: str  # 房源ID\n",
    "    reviewer_id: str  # 评论者ID\n",
    "    reviewer_name: Optional[str] = None  # 评论者姓名，可选\n",
    "    comments: Optional[str] = None  # 评论内容，可选"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a47f72d6-6e88-4de6-b2bc-5c4a60905452",
   "metadata": {
    "height": 744
   },
   "outputs": [],
   "source": [
    "# 定义 Listing 类，用于表示房源相关信息\n",
    "class Listing(BaseModel):\n",
    "    _id: int  # 房源ID\n",
    "    listing_url: str  # 房源URL\n",
    "    name: str  # 房源名称\n",
    "    summary: str  # 房源简介\n",
    "    space: str  # 房源空间描述\n",
    "    description: str  # 房源详细描述\n",
    "    neighborhood_overview: Optional[str] = None  # 邻里概况，可选\n",
    "    notes: Optional[str] = None  # 备注，可选\n",
    "    transit: Optional[str] = None  # 交通信息，可选\n",
    "    access: str  # 访问说明\n",
    "    interaction: Optional[str] = None  # 互动说明，可选\n",
    "    house_rules: str  # 房屋规则\n",
    "    property_type: str  # 房产类型\n",
    "    room_type: str  # 房间类型\n",
    "    bed_type: str  # 床型\n",
    "    minimum_nights: int  # 最少入住天数\n",
    "    maximum_nights: int  # 最多入住天数\n",
    "    cancellation_policy: str  # 取消政策\n",
    "    last_scraped: Optional[datetime] = None  # 上次抓取时间，可选\n",
    "    calendar_last_scraped: Optional[datetime] = None  # 日历上次抓取时间，可选\n",
    "    first_review: Optional[datetime] = None  # 第一条评论时间，可选\n",
    "    last_review: Optional[datetime] = None  # 最后一条评论时间，可选\n",
    "    accommodates: int  # 可容纳人数\n",
    "    bedrooms: Optional[float] = 0  # 卧室数量，可选，默认为0\n",
    "    beds: Optional[float] = 0  # 床位数量，可选，默认为0\n",
    "    number_of_reviews: int  # 评论数量\n",
    "    bathrooms: Optional[float] = 0  # 浴室数量，可选，默认为0\n",
    "    amenities: List[str]  # 设施列表\n",
    "    price: int  # 价格\n",
    "    security_deposit: Optional[float] = None  # 押金金额，可选\n",
    "    cleaning_fee: Optional[float] = None  # 清洁费，可选\n",
    "    extra_people: int  # 额外人数收费\n",
    "    guests_included: int  # 包含的客人数\n",
    "    images: dict  # 图片信息\n",
    "    host: Host  # 房东信息\n",
    "    address: Address  # 地址信息\n",
    "    availability: dict  # 可用性信息\n",
    "    review_scores: dict  # 评论评分\n",
    "    reviews: List[Review]  # 评论列表\n",
    "    text_embeddings: List[float]  # 文本嵌入向量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3f49d2e-0d2d-4445-84da-1bfd3be8b9ff",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": [
    "# 将数据框转换为字典列表形式\n",
    "records = dataset_df.to_dict(orient='records')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8e45574-4e91-4773-9da3-2268939dc72a",
   "metadata": {
    "height": 217
   },
   "outputs": [],
   "source": [
    "# 处理包含 NaT 值的数据\n",
    "for record in records:\n",
    "    for key, value in record.items():\n",
    "        # 检查值是否是列表类型，如果是则处理每个元素\n",
    "        if isinstance(value, list):\n",
    "            processed_list = [None if pd.isnull(v) else v for v in value]\n",
    "            record[key] = processed_list\n",
    "        # 如果是标量值，按之前的方式处理\n",
    "        else:\n",
    "            if pd.isnull(value):\n",
    "                record[key] = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04783976-9ebd-4153-b0ed-cab12505569d",
   "metadata": {
    "height": 149
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    # 将每个字典转换为 Listing 实例\n",
    "    listings = [Listing(**record).dict() for record in records]\n",
    "    # 获取单个数据点的概览\n",
    "    print(listings[0].keys())\n",
    "except ValidationError as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8da84aef-d361-4955-b40c-cf3f05600729",
   "metadata": {},
   "source": [
    "## 1.3 Database Creation and Connection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0d917b2-f08c-4861-9954-c7941d0b8f0c",
   "metadata": {
    "height": 47
   },
   "outputs": [],
   "source": [
    "# 导入所需的库\n",
    "from pymongo.mongo_client import MongoClient\n",
    "from pymongo.operations import SearchIndexModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d076882-9108-4394-8594-f59a12bf2782",
   "metadata": {
    "height": 47
   },
   "outputs": [],
   "source": [
    "# 设置数据库和集合名称\n",
    "database_name = \"airbnb_dataset\"\n",
    "collection_name = \"listings_reviews\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "317d102e-5f93-4444-a7eb-29e0f7954589",
   "metadata": {
    "height": 166
   },
   "outputs": [],
   "source": [
    "def get_mongo_client(mongo_uri: str) -> MongoClient:\n",
    "    \"\"\"建立与 MongoDB 的连接。\n",
    "\n",
    "    Args:\n",
    "        mongo_uri (str): MongoDB 的连接 URI。\n",
    "\n",
    "    Returns:\n",
    "        MongoClient: MongoDB 客户端实例。\n",
    "    \"\"\"\n",
    "    # 与 MongoDB 数据库集群交互的网关\n",
    "    client = MongoClient(mongo_uri, appname=\"devrel.deeplearningai.lesson1.python\")\n",
    "    print(\"Connection to MongoDB successful\")\n",
    "    return client\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abb10772-e0cb-42e3-8b47-71673ec3a130",
   "metadata": {
    "height": 166
   },
   "outputs": [],
   "source": [
    "# 检查 MONGO_URI 是否在环境变量中设置\n",
    "if not MONGO_URI:\n",
    "    print(\"MONGO_URI not set in environment variables\")\n",
    "\n",
    "# 获取 MongoDB 客户端\n",
    "mongo_client = get_mongo_client(MONGO_URI)\n",
    "\n",
    "# 获取指定数据库和集合的 PyMongo 客户端\n",
    "db = mongo_client.get_database(database_name)\n",
    "collection = db.get_collection(collection_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bf7fa64-1abb-4477-9f3d-e7ed5cc6bd34",
   "metadata": {
    "height": 47
   },
   "outputs": [],
   "source": [
    "# 删除集合中所有现有的记录\n",
    "collection.delete_many({})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab5db690-01c1-4439-a346-336859563ab2",
   "metadata": {},
   "source": [
    "## 1.4 Data Ingestion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4f78e05-9297-4f18-8fd1-8c2198b3e5a6",
   "metadata": {
    "height": 64
   },
   "outputs": [],
   "source": [
    "# 数据插入过程可能需要几分钟时间\n",
    "collection.insert_many(listings)\n",
    "print(\"Data ingestion into MongoDB completed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "504a9de1-86d5-48ea-b5d4-9c52eb0424d8",
   "metadata": {},
   "source": [
    "## 1.5 Vector Search Index defintion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb74480e-b469-4d4c-a4e4-a3bdd4acd57a",
   "metadata": {
    "height": 115
   },
   "outputs": [],
   "source": [
    "# 注意：此数据集包含文本和图像嵌入，但本教程仅使用文本嵌入\n",
    "# 每个文档中包含文本嵌入的字段名称\n",
    "text_embedding_field_name = \"text_embeddings\"\n",
    "\n",
    "# MongoDB Atlas 向量搜索索引名称\n",
    "vector_search_index_name_text = \"vector_index_text\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43a30185-7971-4497-b4b4-059dd0e4a970",
   "metadata": {
    "height": 285
   },
   "outputs": [],
   "source": [
    "# 定义向量搜索索引模型\n",
    "vector_search_index_model = SearchIndexModel(\n",
    "    definition={\n",
    "        \"mappings\": {  # 描述数据库文档中字段的索引和存储方式\n",
    "            \"dynamic\": True,  # 自动索引文档中出现的新字段\n",
    "            \"fields\": {  # 将被索引的字段的属性\n",
    "                text_embedding_field_name: { \n",
    "                    \"dimensions\": 1536,  # 向量的维度\n",
    "                    \"similarity\": \"cosine\",  # 计算向量之间相似度的算法\n",
    "                    \"type\": \"knnVector\",  # 向量类型\n",
    "                }\n",
    "            },\n",
    "        }\n",
    "    },\n",
    "    name=vector_search_index_name_text,  # 向量搜索索引的标识符\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "beeb55e8-2ba2-4205-a0da-549e40662e27",
   "metadata": {
    "height": 149
   },
   "outputs": [],
   "source": [
    "# 检查向量搜索索引是否已经存在\n",
    "index_exists = False  # 初始化索引存在标志为 False\n",
    "\n",
    "# 遍历集合中的所有索引\n",
    "for index in collection.list_indexes():\n",
    "    print(index)  # 打印索引信息以供检查\n",
    "    # 如果索引名称与指定的向量搜索索引名称匹配\n",
    "    if index['name'] == vector_search_index_name_text:\n",
    "        index_exists = True  # 设置索引存在标志为 True\n",
    "        break  # 退出循环，因为我们已经找到了目标索引\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dfbcb84-3c10-4ebe-9db1-96c55810a931",
   "metadata": {
    "height": 302
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "# 创建索引（如果索引不存在）\n",
    "if not index_exists:\n",
    "    try:\n",
    "        # 创建搜索索引\n",
    "        result = collection.create_search_index(model=vector_search_index_model)\n",
    "        print(\"Creating index...\")  # 打印正在创建索引的消息\n",
    "        time.sleep(20)  # 休眠20秒，确保向量索引在使用前完成初始同步\n",
    "        print(\"Index created successfully:\", result)  # 打印索引创建成功的消息\n",
    "        print(\"Wait a few minutes before conducting search with index to ensure index initialization\")  # 提示等待几分钟以确保索引初始化完成\n",
    "    except Exception as e:\n",
    "        # 如果创建索引过程中出现异常，打印错误信息\n",
    "        print(f\"Error creating vector search index: {str(e)}\")\n",
    "else:\n",
    "    # 如果索引已存在，打印索引已存在的消息\n",
    "    print(f\"Index '{vector_search_index_name_text}' already exists.\")\n",
    "\n",
    "# 注意：如果此过程的输出为“Error creating vector search index: Duplicate Index”，\n",
    "# 您可以继续下一单元格，如果您打算仍然使用先前创建的索引"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d9134f5-9838-491a-b6e8-0b8271ca218a",
   "metadata": {},
   "source": [
    "<p style=\"background-color:#fff6e4; padding:15px; border-width:3px; border-color:#f5ecda; border-style:solid; border-radius:6px\"> ⏳ <b>Note:</b> If the output of the previous cell is <code>Error creating vector search index: Duplicate Index</code> you may proceed to the next cell if you intend to still use a previously created index.</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "654e68a1-8321-4636-8423-1835ab22b376",
   "metadata": {
    "height": 370
   },
   "outputs": [],
   "source": [
    "import openai\n",
    "\n",
    "# 设置 OpenAI API 密钥\n",
    "openai.api_key = OPENAI_API_KEY\n",
    "\n",
    "def get_embedding(text: str):\n",
    "    \"\"\"使用 OpenAI 的 API 为给定文本生成嵌入向量。\n",
    "    \n",
    "    Args:\n",
    "        text (str): 需要生成嵌入的文本。\n",
    "        \n",
    "    Returns:\n",
    "        list: 生成的嵌入向量，或 None 如果发生错误。\n",
    "    \"\"\"\n",
    "    # 检查输入是否有效\n",
    "    if not text or not isinstance(text, str):  # 如果文本为空或不是字符串类型\n",
    "        return None  # 返回 None\n",
    "\n",
    "    try:\n",
    "        # 调用 OpenAI API 获取嵌入向量\n",
    "        response = openai.Embedding.create(  # 调用 OpenAI 的嵌入向量生成 API\n",
    "            input=text,  # 输入文本\n",
    "            model=\"text-embedding-3-small\"  # 使用的嵌入模型\n",
    "        )\n",
    "        # 从 API 响应中提取嵌入向量\n",
    "        embedding = response['data'][0]['embedding']  # 嵌入向量存储在响应数据的第一条记录中\n",
    "        return embedding  # 返回嵌入向量\n",
    "    except Exception as e:  # 捕捉异常情况\n",
    "        print(f\"Error in get_embedding: {e}\")  # 打印错误信息\n",
    "        return None  # 返回 None 以表示出错"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "246983db-02d2-41a2-a9ae-48ad2f23a4ab",
   "metadata": {},
   "source": [
    "## 1.6 Compose Vector Search Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "884ddc3a-5cc0-47f7-b4cf-c78d7b3b0231",
   "metadata": {
    "height": 931
   },
   "outputs": [],
   "source": [
    "def vector_search(user_query: str, db, collection, vector_index=\"vector_index_text\"):\n",
    "    \"\"\"\n",
    "    在 MongoDB 集合中基于用户查询执行向量搜索。\n",
    "\n",
    "    Args:\n",
    "    user_query (str): 用户的查询字符串。\n",
    "    db (MongoClient.database): 数据库对象。\n",
    "    collection (MongoCollection): 要搜索的 MongoDB 集合。\n",
    "    vector_index (str): 向量索引名称，默认为 \"vector_index_text\"。\n",
    "\n",
    "    Returns:\n",
    "    list: 匹配的文档列表。\n",
    "    \"\"\"\n",
    "\n",
    "    # 为用户查询生成嵌入向量\n",
    "    query_embedding = get_embedding(user_query)\n",
    "\n",
    "    if query_embedding is None:\n",
    "        return \"Invalid query or embedding generation failed.\"\n",
    "\n",
    "    # 定义向量搜索阶段\n",
    "    vector_search_stage = {\n",
    "        \"$vectorSearch\": {\n",
    "            \"index\": vector_index,  # 指定用于搜索的索引\n",
    "            \"queryVector\": query_embedding,  # 表示查询的向量\n",
    "            \"path\": text_embedding_field_name,  # 文档中包含要搜索向量的字段\n",
    "            \"numCandidates\": 150,  # 考虑的候选匹配数\n",
    "            \"limit\": 20  # 返回前20个匹配结果\n",
    "        }\n",
    "    }\n",
    "\n",
    "    # 定义包含向量搜索阶段和其他阶段的聚合管道\n",
    "    pipeline = [vector_search_stage]\n",
    "\n",
    "    # 执行搜索\n",
    "    results = collection.aggregate(pipeline)\n",
    "\n",
    "    # 解释查询执行计划\n",
    "    explain_query_execution = db.command(\n",
    "        'explain', {  # 返回有关 MongoDB 如何执行查询或命令的信息，而无需实际运行它\n",
    "            'aggregate': collection.name,  # 指定执行聚合的集合名称\n",
    "            'pipeline': pipeline,  # 要分析的聚合管道\n",
    "            'cursor': {}  # 指示应使用默认游标行为\n",
    "        }, \n",
    "        verbosity='executionStats'  # 有关聚合管道每个阶段执行的详细统计信息\n",
    "    )\n",
    "\n",
    "    # 获取向量搜索的解释信息\n",
    "    vector_search_explain = explain_query_execution['stages'][0]['$vectorSearch']\n",
    "    millis_elapsed = vector_search_explain['explain']['collectStats']['millisElapsed']\n",
    "\n",
    "    # 打印数据库服务器上执行完成所需的总时间\n",
    "    print(f\"Total time for the execution to complete on the database server: {millis_elapsed} milliseconds\")\n",
    "\n",
    "    return list(results)  # 返回搜索结果列表\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa2107cd-c922-492a-ad4f-54cb94b436a3",
   "metadata": {},
   "source": [
    "## 1.7 Handling User Query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b771954d-e823-42b7-9141-5f74298b332e",
   "metadata": {
    "height": 149
   },
   "outputs": [],
   "source": [
    "# 定义 SearchResultItem 类，用于表示搜索结果项的相关信息\n",
    "class SearchResultItem(BaseModel):\n",
    "    name: str  # 房源名称\n",
    "    accommodates: Optional[int] = None  # 可容纳人数，可选\n",
    "    address: Address  # 地址信息\n",
    "    summary: Optional[str] = None  # 简介，可选\n",
    "    description: Optional[str] = None  # 详细描述，可选\n",
    "    neighborhood_overview: Optional[str] = None  # 邻里概况，可选\n",
    "    notes: Optional[str] = None  # 备注，可选"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49ff7a10-0ae4-4e7d-9e30-5a144e742441",
   "metadata": {
    "height": 778
   },
   "outputs": [],
   "source": [
    "from IPython.display import display, HTML\n",
    "\n",
    "def handle_user_query(query, db, collection):\n",
    "    \"\"\"\n",
    "    处理用户查询并返回系统响应和源信息。\n",
    "\n",
    "    Args:\n",
    "    query (str): 用户的查询字符串。\n",
    "    db (MongoClient.database): 数据库对象。\n",
    "    collection (MongoCollection): 要搜索的 MongoDB 集合。\n",
    "\n",
    "    Returns:\n",
    "    str: 系统响应。\n",
    "    \"\"\"\n",
    "    # 执行向量搜索，假设返回的是一个包含 'title' 和 'plot' 键的字典列表\n",
    "    get_knowledge = vector_search(query, db, collection)\n",
    "\n",
    "    # 检查是否有结果\n",
    "    if not get_knowledge:\n",
    "        return \"No results found.\", \"No source information available.\"\n",
    "        \n",
    "    # 将搜索结果转换为 SearchResultItem 模型列表\n",
    "    search_results_models = [\n",
    "        SearchResultItem(**result)\n",
    "        for result in get_knowledge\n",
    "    ]\n",
    "\n",
    "    # 将搜索结果转换为 DataFrame 以便在 Jupyter 中更好地呈现\n",
    "    search_results_df = pd.DataFrame([item.dict() for item in search_results_models])\n",
    "\n",
    "    # 使用 OpenAI 的 completion 生成系统响应\n",
    "    completion = openai.ChatCompletion.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"system\", \n",
    "                \"content\": \"You are an Airbnb listing recommendation system.\"\n",
    "            },\n",
    "            {\n",
    "                \"role\": \"user\", \n",
    "                \"content\": f\"Answer this user query: {query} with the following context:\\n{search_results_df}\"\n",
    "            }\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    system_response = completion.choices[0].message['content']\n",
    "\n",
    "    # 打印用户问题、系统响应和源信息\n",
    "    print(f\"- User Question:\\n{query}\\n\")\n",
    "    print(f\"- System Response:\\n{system_response}\\n\")\n",
    "\n",
    "    # 以 HTML 表格形式显示 DataFrame\n",
    "    display(HTML(search_results_df.to_html()))\n",
    "\n",
    "    # 返回结构化响应和源信息作为字符串\n",
    "    return system_response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41b800f9-b005-4800-af73-ed36c552e7ad",
   "metadata": {
    "height": 132
   },
   "outputs": [],
   "source": [
    "query = \"\"\"\n",
    "I want to stay in a place that's warm and friendly, \n",
    "and not too far from resturants, can you recommend a place? \n",
    "Include a reason as to why you've chosen your selection.\n",
    "\"\"\"\n",
    "handle_user_query(query, db, collection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2021f9a3-0b67-493c-a5e3-d8ea2a695e63",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5502d2e5-c125-4141-8dae-df413c6f050d",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95a5ea3b-ca85-494d-a7b1-abc1b0584972",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a33a1f4-459a-4c49-b39e-04fb8268ab26",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00471809-c526-4e58-833e-29c9b449ce21",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b089114-f92c-46b8-8b51-cb218d7a8ff9",
   "metadata": {
    "height": 30
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
